import os
import math
import json
from PIL import Image
import argparse

try:
    import numpy as np
except ImportError:
    raise ImportError('请先安装 numpy: pip install numpy')
# 生成标准Web Mercator坐标系的瓦片(采用z/x/y.png编号格式，x/y范围与主流地图一致)，可使用任何标准地图前端(如index.html)直接访问这些瓦片
def mercator_to_pixel(mx, my, zoom, tile_size=512):
    initial_resolution = 2 * math.pi * 6378137.0 / tile_size
    origin_shift = 2 * math.pi * 6378137.0 / 2.0
    res = initial_resolution / (2 ** zoom)
    px = (mx + origin_shift) / res
    py = (origin_shift - my) / res
    return px, py

# 优先使用config.json中精确的web_mercator坐标
def mercator_to_pixel_from_config(mx, my, zoom, tile_size=512):
    return mercator_to_pixel(mx, my, zoom, tile_size)

# 从配置文件加载参数
with open('config.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

# 使用配置中的参数
corners_wgs84 = config['corners']['wgs84']
IMG_WIDTH = config['img_width']
IMG_HEIGHT = config['img_height']

def generate_tiles(input_image=None, output_dir=None, zoom_min=None, zoom_max=None, tile_size=None):
    # 使用配置参数作为默认值
    input_image = input_image or config['input_image']
    output_dir = output_dir or config['output_dir']
    zoom_min = zoom_min or config['min_zoom']
    zoom_max = zoom_max or config['max_zoom']
    tile_size = tile_size or config['tile_size']
    img = Image.open(input_image)
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    img_width, img_height = img.size
    print(f"输入图片尺寸: {img_width}x{img_height}")

    # 直接用config.json的web_mercator四角和原始图片，不再依赖gdalwarp

    # 计算图片四角在像素坐标系下的位置
    for z in range(zoom_min, zoom_max + 1):
        print(f"正在处理缩放级别 {z}...")
        # 用config.json中web_mercator四角严格计算应生成的瓦片范围
        px_corners = [mercator_to_pixel_from_config(mx, my, z, tile_size) for mx, my in config['corners']['web_mercator']]
        pxs = [p[0] for p in px_corners]
        pys = [p[1] for p in px_corners]
        min_px, max_px = min(pxs), max(pxs)
        min_py, max_py = min(pys), max(pys)
        # 用标准Web Mercator tile算法，向外取整，保证所有覆盖范围内的tile都生成
        x_tile_min = int(math.floor(min_px / tile_size))
        x_tile_max = int(math.ceil(max_px / tile_size)) - 1
        y_tile_min = int(math.floor(min_py / tile_size))
        y_tile_max = int(math.ceil(max_py / tile_size)) - 1
        print(f"瓦片坐标范围: x {x_tile_min}-{x_tile_max}, y {y_tile_min}-{y_tile_max}")

        # 仿射变换严格用config.json四角和图片像素四角一一对应
        src = [
            (0, img_height - 1),           # 左下
            (0, 0),                        # 左上
            (img_width - 1, 0),            # 右上
            (img_width - 1, img_height - 1) # 右下
        ]
        dst = px_corners
        def get_affine_transform(src, dst):
            A = []
            B = []
            for i in range(4):
                x, y = src[i]
                u, v = dst[i]
                A.append([x, y, 1, 0, 0, 0])
                A.append([0, 0, 0, x, y, 1])
                B.append(u)
                B.append(v)
            A = np.array(A)
            B = np.array(B)
            affine, _, _, _ = np.linalg.lstsq(A, B, rcond=None)
            return affine.reshape(2, 3)
        affine = get_affine_transform(src, dst)
        def geo_to_img(px, py):
            A = affine
            A_inv = np.linalg.pinv(np.vstack([A, [0,0,1]]))
            vec = np.array([px, py, 1])
            img_xy = A_inv @ vec
            return img_xy[0], img_xy[1]

        # 计算当前zoom下地理像素分辨率（即图片在当前zoom下应有的像素宽高）
        if z >= 17:
            # 计算地理像素范围
            px_corners = [mercator_to_pixel_from_config(mx, my, z, tile_size) for mx, my in config['corners']['web_mercator']]
            pxs = [p[0] for p in px_corners]
            pys = [p[1] for p in px_corners]
            min_px, max_px = min(pxs), max(pxs)
            min_py, max_py = min(pys), max(pys)
            geo_width = int(np.ceil(max_px - min_px))
            geo_height = int(np.ceil(max_py - min_py))
            # 将整张图片resize到地理像素分辨率
            big_img = img.resize((geo_width, geo_height), Image.Resampling.LANCZOS)
        for x in range(x_tile_min, x_tile_max + 1):
            for y in range(y_tile_min, y_tile_max + 1):
                px0, py0 = x * tile_size, y * tile_size
                px1, py1 = (x+1) * tile_size, (y+1) * tile_size
                if z >= 17:
                    # 在放大后的图片上切tile
                    # 计算当前tile在big_img上的像素范围
                    left = int(np.floor(px0 - min_px))
                    upper = int(np.floor(py0 - min_py))
                    right = int(np.ceil(px1 - min_px))
                    lower = int(np.ceil(py1 - min_py))
                    # 边界裁剪
                    left = max(left, 0)
                    upper = max(upper, 0)
                    right = min(right, big_img.width)
                    lower = min(lower, big_img.height)
                    if left < right and upper < lower:
                        region = big_img.crop((left, upper, right, lower))
                        region = region.resize((tile_size, tile_size), Image.Resampling.LANCZOS)
                        tile_img = Image.new('RGBA', (tile_size, tile_size), (0,0,0,0))
                        tile_img.paste(region, (0,0))
                    else:
                        tile_img = Image.new('RGBA', (tile_size, tile_size), (0,0,0,0))
                else:
                    # 原有低zoom逻辑
                    img_pts = [geo_to_img(px0, py0), geo_to_img(px1, py0), geo_to_img(px1, py1), geo_to_img(px0, py1)]
                    xs = [p[0] for p in img_pts]
                    ys = [p[1] for p in img_pts]
                    min_x = int(np.floor(min(xs)))
                    max_x = int(np.ceil(max(xs)))
                    min_y = int(np.floor(min(ys)))
                    max_y = int(np.ceil(max(ys)))
                    region_left = max(min_x, 0)
                    region_upper = max(min_y, 0)
                    region_right = min(max_x, img_width)
                    region_lower = min(max_y, img_height)
                    if region_left < region_right and region_upper < region_lower:
                        paste_x = region_left - min_x
                        paste_y = region_upper - min_y
                        region = img.crop((region_left, region_upper, region_right, region_lower))
                        region_width = region_right - region_left
                        region_height = region_lower - region_upper
                        scale_x = tile_size / (max_x - min_x) if (max_x - min_x) != 0 else tile_size
                        scale_y = tile_size / (max_y - min_y) if (max_y - min_y) != 0 else tile_size
                        region = region.resize((int(region_width * scale_x), int(region_height * scale_y)), Image.Resampling.LANCZOS)
                        tile_img = Image.new('RGBA', (tile_size, tile_size), (0,0,0,0))
                        tile_img.paste(region, (int(paste_x * scale_x), int(paste_y * scale_y)))
                    else:
                        tile_img = Image.new('RGBA', (tile_size, tile_size), (0,0,0,0))
                tile_dir = os.path.join(output_dir, str(z))
                os.makedirs(tile_dir, exist_ok=True)
                tile_path = os.path.join(tile_dir, f"{x}_{y}.png")
                tile_img.save(tile_path)
        print(f"缩放级别 {z} 处理完成，瓦片已保存至 {os.path.join(output_dir, str(z))}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='从地理参考图像生成标准Web Mercator瓦片')
    parser.add_argument('--input', type=str, default=None, help='输入图片文件(默认从config.json读取)')
    parser.add_argument('--output', type=str, default=None, help='瓦片输出目录(默认从config.json读取)')
    parser.add_argument('--minzoom', type=int, default=None, help='最小缩放级别(默认从config.json读取)')
    parser.add_argument('--maxzoom', type=int, default=None, help='最大缩放级别(默认从config.json读取)')
    parser.add_argument('--tile_size', type=int, default=None, help='瓦片尺寸(默认从config.json读取)')
    args = parser.parse_args()

    print(f"开始生成瓦片...")
    if args.input is None and args.output is None and args.minzoom is None and args.maxzoom is None and args.tile_size is None:
        print("未指定命令行参数，已从config.json读取参数：")
        print(f"输入图片: {config['input_image']}")
        print(f"输出目录: {config['output_dir']}")
        print(f"缩放级别范围: {config['min_zoom']} 至 {config['max_zoom']}")
        print(f"瓦片尺寸: {config['tile_size']}")
    else:
        print(f"输入图片: {args.input}")
        print(f"输出目录: {args.output}")
        print(f"缩放级别范围: {args.minzoom} 至 {args.maxzoom}")
        print(f"瓦片尺寸: {args.tile_size}")

    generate_tiles(
        input_image=args.input,
        output_dir=args.output,
        zoom_min=args.minzoom,
        zoom_max=args.maxzoom,
        tile_size=args.tile_size
    )
    print("瓦片生成完成!")